# coding: utf-8

from __future__ import print_function
import jieba

import tensorflow as tf
import tensorflow.contrib.keras as kr

from src.classification.material.textcnn.cnn_model import TCNNConfig, TextCNN
from src.classification.material.textcnn.data_loader import read_category, read_vocab
from configs.classification_config import MATERIAL_VOCAL_PATH as VOCAL_PATH, MATERIAL_TEXTCNN_MODEL_PATH as TEXTCNN_MODEL_PATH


class CnnModel:
    def __init__(self):
        self.config = TCNNConfig()
        self.categories, self.cat_to_id = read_category()
        self.words, self.word_to_id = read_vocab(VOCAL_PATH)
        self.config.vocab_size = len(self.words)
        tf.reset_default_graph()
        self.model = TextCNN(self.config)

        self.session = tf.Session()
        self.session.run(tf.global_variables_initializer())
        saver = tf.train.Saver()
        saver.restore(sess=self.session, save_path=TEXTCNN_MODEL_PATH)  # 读取保存的模型

    def predict(self, messages, topK=1):
        # content = message.strip()
        data = []
        for content in messages:
            words = jieba.lcut(content)
            data.append([self.word_to_id[x] for x in words if x in self.word_to_id])

        feed_dict = {
            self.model.input_x: kr.preprocessing.sequence.pad_sequences(data, self.config.seq_length),
            self.model.keep_prob: 1.0
        }

        top_y_pred_cls = self.session.run(self.model.top_y_pred_cls, feed_dict=feed_dict)

        # return self.categories[y_pred_cls[0]]
        res = []
        for item in top_y_pred_cls:
            temp = []
            for i, p in enumerate(item):
                temp.append((self.categories[i], p))
            temp.sort(key=lambda x:x[1], reverse=True)
            res.append(temp[:topK])

        return res
